import torch
from torch import nn

# ========================加载模型结构=============================
from torchvision.models import resnet18, ResNet18_Weights

# 指定训练的设备
device = torch.device("cuda")

model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

# =========== 如果用GPU去训练模型,则模型默认的设备是GPU
model = model.to(device)
model.fc = nn.Linear(512, 3,device=device)
# ===========================加载权重 ==========================
state_dict = torch.load("./save/best1.pt")  # 主要问题就是此处权重是GPU
model.load_state_dict(state_dict)  # 模型结构加入权重

# ========================加载数据=============================
import cv2
import numpy as np

image = cv2.imread("./data/valid/apple/50.jpg")  # numpy数据
image = cv2.resize(image, (224, 224))
image = np.expand_dims(image, 0)
image = torch.from_numpy(image).permute([0, 3, 1, 2])

model.eval()  # 启动验证模式
result = model(image.to(device).float())  # 预测结果（此时结果为softmax值）,输入值也需要GPU
result = torch.argmax(result, dim=-1)  # 得到最终的预测的数字内容（此处因为下标与数字的值是一致的，因此不需要映射表）
print(result)
